# Load Lasso data to get optimal lambda
for (mod in c("lasso","rf","nn","gbm")) {
  load(paste(modeldir,"/",
             country,
             "_",mod,"_",
             "demean",
             "_",
             rhs.group,fileext,
             ".RData",
             sep = ""))
}
#####################################
######### B- GET ESTIMATES ##########
#####################################
#####################################

# Set holder for each year's results/model
ebma.results <- list()
# Loop over years
i = 0
for (year in start.year:end.year) { 
  i = i + 1
  print(year)
  # Generate dependent variable
  Violence.mean <-  aggregate(x = as.matrix(dta[dta[,t.var]< year,
                                                     dv]),
                              by = list(dta[dta[,t.var]< year,
                                                 loc.var]),
                              FUN = mean)
  colnames(Violence.mean) <- c(loc.var,
                               "dv_mean")
  Violence.new <- merge(as.matrix(dta[dta[,t.var]< year,
                                           c(loc.var,t.var,dv)]),
                        Violence.mean,
                        by = loc.var)
  Violence.past <- Violence.new[,dv]-Violence.new$dv_mean
  dta.past <- dta[dta[,t.var]< year,]
  N <- length(dta.past[,dv])
  ###################################
  ### 1 # Find Optimal Parameters ###
  ###################################
  
  # Do repeated cross validation to find best weights and DT
  cv_results <- foreach (cv=1:cv.runs, 
                         .combine='rbind',
                         .errorhandling = 'remove',
                         .packages = c("glmnet",
                                       "randomForest",
                                       "gbm",
                                       "nnet",
                                       "pROC")) %dopar% {   # Loop over CV runs
                                         shuffle <- sample(1:N,
                                                           N,
                                                           replace=FALSE)
                                         sampleSplit = c(0,round((N/5*c(1:5))))
    # Set matrix to store predictions
    pred <- matrix(NA,nrow=N,ncol=5)
    colnames(pred) <- c("lasso.prediction",
                        "rf.prediction",
                        "gbm.prediction",
                        "nn.prediction",
                        "ebma.prediction")
    for (f in 1:5) {     # Loop over folds       
      # Fit lasso
      lasso.fit <-  glmnet(x = as.matrix(dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                  rhs]),
                           y = as.matrix(Violence.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])]]),
                           alpha = lasso_demean_params$alpha,
                           family = lasso_demean_params$family,
                           lambda = seq(10.9,1,-.1)*lasso.results[[i]]$lambda)
      opt.lam <- 100
      while (is.na(lasso.fit$lambda[opt.lam])) {
        opt.lam = opt.lam - 1
      }
      # Save predictions for fold f
      pred[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
           "lasso.prediction"] <- predict(lasso.fit,
                                                          newx = as.matrix(dta.past[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                                                         rhs]),
                                                          type = "response")[,opt.lam]
      # Fit Random Forest
      rf.fit <- randomForest(x = as.matrix(dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                         rhs]),
                             y = as.matrix(Violence.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])]]),
                             ntree = rf_demean_params$ntree,
                             nodesize = rf_demean_params$nodesize,
                             mtry=rf_demean_params$mtry,
                             type = rf_demean_params$type,
                             do.trace = rf_demean_params$do.trace,
                             importance = rf_demean_params$importance
      )
      # Save predictions for fold f
      pred[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],"rf.prediction"] <- predict(rf.fit,
                                                       newdata = as.matrix(dta.past[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                                                         rhs]))
      # Fit AdaBoost
      gbm.fit <- gbm.fit(x = as.matrix(dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                     rhs]),
                         y = as.matrix(Violence.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])]]),
                         n.trees = gbm.results[[i]]$trees,
                         distribution = gbm_demean_params$distribution,
                         interaction.depth = gbm_demean_params$interaction.depth,
                         n.minobsinnode = gbm_demean_params$n.minobsinnode,
                         verbose = gbm_demean_params$verbose,
                         shrinkage = gbm_demean_params$shrinkage,
                         train.fraction = gbm_demean_params$train.fraction,
                         bag.fraction = gbm_demean_params$bag.fraction)
      # Save predictions for fold f
      pred[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],"gbm.prediction"] <- predict(gbm.fit,
                                                        newdata = as.matrix(dta.past[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                                                                          rhs]),
                                                        type = "response",
                                                        n.trees = gbm.results[[i]]$trees)
      
      # Fit Neural Network
      # Remove constant variables from predictors
      train.RHS <- dta.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                 rhs]
      test.RHS <- dta.past[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],
                                rhs]
      # Remove variables that are constant
      not.constant <- apply(train.RHS,2,sd)!=0
      train.RHS <- train.RHS[,not.constant]
      test.RHS <- test.RHS[,not.constant]
      # Find PC weights
      pcs.fit <- prcomp(train.RHS,center = TRUE, scale = TRUE)
      train.pcs <- predict(pcs.fit,
                           train.RHS)[,1:min(nn_demean_params$pcNum,
                                             ncol(train.RHS))]
      test.pcs <- predict(pcs.fit,
                          test.RHS)[,1:min(nn_demean_params$pcNum,
                                           ncol(train.RHS))]
      # Fit Network
      nn.fit <- nnet(x = train.pcs,
                     y = as.matrix(Violence.past[-shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])]]),
                     size = nn.results[[i]]$size,
                     decay = nn_demean_params$decay,
                     trace = nn_demean_params$trace,
                     linout = nn_demean_params$linout
      )
      
      # Save predictions for fold f
      pred[shuffle[(sampleSplit[f]+1):(sampleSplit[f+1])],"nn.prediction"] <- predict(nn.fit,
                                                       newdata = test.pcs)
      
    }
    # Get weights for each model
    f <- pred[,c("lasso.prediction",
                      "rf.prediction",
                      "gbm.prediction",
                      "nn.prediction")]
    g <- f*0
    sigma <- colMeans((matrix(rep(Violence.past,4),
                              byrow = FALSE,
                              ncol = 4)-f)^2)
    max.dif <- 1
    while (max.dif >.01) {
      for (col in 1:4) {
        for (row  in 1:nrow(g))
        g[row,col] <- pnorm(Violence.past[row],
                   mean = f[row,col],
                   sd = sqrt(sigma[col]))
      }
      z <- g/matrix(rep(rowSums(g),4),
                    byrow = FALSE,
                    ncol = 4)
      w <- colMeans(z)
      sigma.new <- colSums(z*(matrix(rep(Violence.past,4),
                                    byrow = FALSE,
                                    ncol = 4)-f)^2)/colSums(z)
      max.dif <- max(abs(sigma-sigma.new))
      sigma <- sigma.new
    }
    pred[,"ebma.prediction"] <- as.matrix(pred[,c("lasso.prediction",
                                                            "rf.prediction",
                                                            "gbm.prediction",
                                                            "nn.prediction")]) %*% as.matrix(w)
    cv_results <- c(cv,w)
    names(cv_results) <- c("CV Run",
                           "lasso.weight",
                           "rf.weight",
                           "gbm.weight",
                           "nn.weight") 
    return(cv_results)
  } # End cv runs
  hours <- floor((proc.time()[3]-start.time)/3600)
  mins <- floor((proc.time()[3]-start.time)/60) - hours*60
  secs <- floor((proc.time()[3]-start.time)) - hours*3600 - mins*60
  print(paste(hours,
              "h",
              mins,
              "m",
              secs,
              "s elapsed",
              sep = ""))
  
  ######################################
  ### 2 # FIT OPTIMAL MODEL ############
  ######################################
  
  # A1 - Fit Model
  ###############
  
  ebma.results[[i]] <- list()
  ebma.results[[i]]$optim.wts <- colMeans(cv_results[,c("lasso.weight",
                                                    "rf.weight",
                                                    "gbm.weight",
                                                    "nn.weight")])
  # B1 Fit out of sample
  ########################
  
  ebma.results[[i]]$outsample.component.prob <- cbind("lasso" = lasso.results[[i]]$fit.oos,
                                                  "rf" = rf.results[[i]]$fit.oos,
                                                  "gbm" = gbm.results[[i]]$fit.oos,
                                                  "nn" = nn.results[[i]]$fit.oos)
  ebma.results[[i]]$fit.oos <- ebma.results[[i]]$outsample.component.prob %*% as.matrix(ebma.results[[i]]$optim.wts)
  Violence.new <- merge(as.matrix(dta[dta[,t.var]== year,
                                           c(loc.var,t.var,dv)]),
                        Violence.mean,
                        by = loc.var)
  ebma.results[[i]]$actual.oos <- Violence.new[,dv]-Violence.new$dv_mean
    
}
save(ebma.results,file=paste(modeldir,"/",
                           country,
                           "_ebma_",
                           "demean",
                           "_",
                           rhs.group,fileext,
                           ".RData",
                           sep = ""))